package class02

import (
	"fmt"
	"math"
)

/*
	如果给你一个二维数组，每个值代表这一块地形的高度
	求整块地形可以装下多少水？
*/

type HeapNode struct {
	Value int
	Row   int
	Col   int
}

func NewHeap(compare func(interface{}, interface{})bool) *MyHeap {
	return &MyHeap{
		compare: compare,
		Nodes:[]*HeapNode{},
		Size: 0,
	}
}

type MyHeap struct {
	compare func(interface{}, interface{})bool
	Nodes []*HeapNode
	Size  int
}

func (h *MyHeap) Insert(n *HeapNode){
	h.Nodes = append(h.Nodes, n)
	h.Up(h.Size)
	h.Size++
}

func (h *MyHeap) Down(index int) {
	cur := index
	small := 2*index + 1

	for small < h.Size {
		if small + 1 < h.Size && !h.compare(h.Nodes[small], h.Nodes[small+1]) {
			small = small + 1
		}
		if !h.compare(h.Nodes[cur], h.Nodes[small]) {
			h.Swap(cur, small)
		} else {
			fmt.Println(small)
			break
		}
		cur = small
		small = 2*cur + 1
	}
}

func (h *MyHeap) Up(index int) {
	parent := (index-1)/2
	cur := index
	for cur > 0 {
		if h.compare(h.Nodes[cur], h.Nodes[parent]) {
			h.Swap(cur, parent)
		} else {
			break
		}
		cur = parent
		parent = (cur-1)/2
	}
}

func (h *MyHeap) Swap(index1, index2 int) {
	h.Nodes[index1], h.Nodes[index2] = h.Nodes[index2], h.Nodes[index1]
}

func (h *MyHeap) Poll() *HeapNode {
	if h.Size == 0 {
		panic(any("not element"))
	}
	h.Size--
	ans := h.Nodes[0]
	h.Swap(0, h.Size)
	h.Nodes = h.Nodes[:len(h.Nodes)-1]
	h.Down(0)
	return ans
}

func (h *MyHeap) IsEmpty() bool {
	return h.Size == 0
}

func TrapRainWater(arr [][]int) int {
	var ans int

	// 创建小根堆
	com := func(n1 interface{}, n2 interface{})bool {
		return n1.(*HeapNode).Value < n2.(*HeapNode).Value
	}
	myH := NewHeap(com)


	// 收集矩阵四周
	le := len(arr[0])
	wi := len(arr)
	// 记录进入过小根堆的位置
	xyRecord := make([][]int, wi)
	for i:=0; i<wi; i++ {
		xyRecord[i] = make([]int, le)
	}
	for i:=0; i<le-1; i++ {
		myH.Insert(&HeapNode{Value: arr[0][i], Col: i, Row: 0})
		xyRecord[0][i] = 1
	}
	for i:=0; i<wi-1; i++ {
		myH.Insert(&HeapNode{Value: arr[i][le-1], Col: le-1, Row: i})
		xyRecord[i][le-1] = 1
	}
	for i:=le-1; i>0; i-- {
		myH.Insert(&HeapNode{Value: arr[wi-1][i], Col: i, Row: wi-1})
		xyRecord[wi-1][i] = 1
	}
	for i:=wi-1; i>0; i-- {
		myH.Insert(&HeapNode{Value: arr[i][0], Col: 0, Row: i})
		xyRecord[i][0] = 1
	}

	maxH := math.MinInt
	for !myH.IsEmpty() {
		minN := myH.Poll()
		maxH = max(maxH, minN.Value)
		// 上方
		if minN.Row - 1 >= 0 && xyRecord[minN.Row-1][minN.Col] != 1 {
			xyRecord[minN.Row-1][minN.Col] = 1
			myH.Insert(&HeapNode{Value: arr[minN.Row-1][minN.Col], Row: minN.Row-1, Col: minN.Col})
			ans += max(0, maxH - arr[minN.Row-1][minN.Col])
		}
		// 下
		if minN.Row + 1 < wi && xyRecord[minN.Row + 1][minN.Col] != 1 {
			xyRecord[minN.Row + 1][minN.Col] = 1
			myH.Insert(&HeapNode{Value: arr[minN.Row + 1][minN.Col], Row: minN.Row + 1, Col: minN.Col})
			ans += max(0, maxH - arr[minN.Row + 1][minN.Col])
		}
		// 左
		if minN.Col - 1 >= 0 && xyRecord[minN.Row][minN.Col-1] != 1 {
			xyRecord[minN.Row][minN.Col-1] = 1
			myH.Insert(&HeapNode{Value: arr[minN.Row][minN.Col-1], Row: minN.Row, Col: minN.Col-1})
			ans += max(0, maxH - arr[minN.Row][minN.Col-1])
		}
		// 右边
		if minN.Col + 1 < le && xyRecord[minN.Row][minN.Col + 1] != 1 {
			xyRecord[minN.Row][minN.Col + 1] = 1
			myH.Insert(&HeapNode{Value: arr[minN.Row][minN.Col + 1], Row: minN.Row, Col: minN.Col + 1})
			ans += max(0, maxH - arr[minN.Row][minN.Col + 1])
		}
	}

	return ans
}
